import base64
import datetime
import json
import os
import sys
from io import BytesIO

import cv2
import insightface
import numpy as np
from flask import Flask, request, jsonify, abort
from gevent.pywsgi import WSGIServer


class FaceAnalysis:
    def __init__(self, gpu_num=-1, nms=0.4):
        self.model = insightface.app.FaceAnalysis()
        self.model.prepare(ctx_id=gpu_num)#, det_thresh=nms)#, nms=nms)

    def __call__(self, x):
        """
        :param x: RGB [m, n, 3] for cv2.imread, not RGB.
        :return:
        """
        faces = self.model.get(x)
        print(len(faces))
        for idx, face in enumerate(faces):
            print("Face [%d]:" % idx)
            print("\tage:%d" % (face.age))
            gender = 'Male'
            if face.gender == 0:
                gender = 'Female'
            print("\tgender:%s" % (gender))
            print("\tembedding shape:%s" % face.embedding.shape)
            print("\tbbox:%s" % (face.bbox.astype(np.int).flatten()))
            #print("\tlandmark:%s" % (face.landmark.astype(np.int).flatten()))
            print("")
        assert len(faces) == 1, "This image must only have one face "

        return {'age': face.age, 'gender': gender, 'bbox':(face.bbox.astype(np.int).flatten()).tolist(),
                'feature_vector': face.embedding.tolist()}


app = Flask(__name__)
# 常量

blackmanba = FaceAnalysis(int(sys.argv[1]))


@app.route('/faceanalysis', methods=['POST'])
def requests():
    data = request.get_json(force=True)
    if (not data or not 'reqid' in data) or (not 'reqtime' in data) or (not 'image' in data) > 1:
        abort(400)

    pic = BytesIO()
    pic.write(base64.b64decode(bytes(data["image"], encoding='utf8')))
    pic.seek(0)
    file_bytes = np.asarray(bytearray(pic.read()), dtype=np.uint8)

    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)[..., ::-1]

    raw_result = blackmanba(img)

    result = {
        'retid': data['reqid'],
        'rettime': datetime.datetime.now().strftime('%Y%m%d%H%M%S'),
    }
    result.update(raw_result)
    result = json.dumps(result)
    return jsonify(result)


if __name__ == '__main__':
    ip = '127.0.0.1'
    post = 5000
    print('running server http://{0}'.format(ip + ':' + str(post)))
    WSGIServer((ip, post), app).serve_forever()
